#coding=utf-8
import matplotlib.pyplot as plt
import numpy as np
import time
import cPickle as pickle
from fakekeras import Sequential  # 顺序模型
from fakekeras import FullConnectedLayer,ConvLayer,MaxPoolingLayer,FlattenLayer #层
from fakekeras import ReluActivator,SoftmaxActivator # 激活函数


# 加载minist字符集
def load_minst(data_file):
    import gzip
    f = gzip.open(data_file, "rb")
    train, val, test = pickle.load(f)
    f.close()
    # 训练数据（图片）的格式 [samples][channels][width][height]
    train_x = train[0].reshape((train[0].shape[0],1,28,28)).astype('float32')
    test_x = test[0].reshape((test[0].shape[0],1,28,28)).astype('float32')
    # 目标数据的格式 [n维目标向量，向量数]  (对于手写字符，n=10)
    train_y = np.zeros((train[1].shape[0],10)).astype('float32')
    test_y = np.zeros((test[1].shape[0],10)).astype('float32')
    for i in xrange(train[1].shape[0]):
        train_y[i][train[1][i]]=1.0
    for i in xrange(test[1].shape[0]):
        test_y[i][test[1][i]]=1.0
    return train_x, train_y.T, test_x, test_y.T

if __name__ == '__main__':
    print 'runing...'
    
    #加载手写数字集 http://deeplearning.net/data/mnist/mnist.pkl.gz
    X_train, y_train, X_test, y_test = load_minst("mnist.pkl.gz")
    
    # 搭建一个网络
    network = Sequential()
    network.add(ConvLayer(3, 3,6,0,1,ReluActivator(),0.1,input_shape=(1,28,28))) # 卷积层
    network.add(MaxPoolingLayer(2,2,2)) # max polling 层
    network.add(ConvLayer(4, 4,13,0,1,ReluActivator(),0.1,input_shape=(6,13,13))) # 卷积层
    network.add(MaxPoolingLayer(2,2,2)) # max polling 层
    network.add(FlattenLayer()) # flatten 层
    network.add(FullConnectedLayer(units=20,activator=ReluActivator()))  # 三个全连接层
    network.add(FullConnectedLayer(units=20,activator=ReluActivator()))
    network.add(FullConnectedLayer(units=10,activator=SoftmaxActivator()))    
    # 训练
    n=6000 #选n个样本来训练
    trained=True
    #trained=True自动加载训练好的的模型
    if trained:
        network.load_model(modelfile='model.npy')
    else:
        network.train(X_train[0:n],y_train[:,0:n],batch_size=100,epochs=20,rate=0.001)  
    # 计算正确率
    print u"正确率 =%f"%(100-100*network.evaluate(X_test,y_test))
    
    # 可视化
    num_classes=10
    np.random.seed(int(time.time()))
    fig = plt.figure() #figsize=(8,3)
    for i in range(num_classes):
        ax = fig.add_subplot(2, 5, 1 + i, xticks=[], yticks=[])
        idx=np.random.randint(0,10000) # 随机取一个
        rs=(network.predict(X_test[idx:idx+1]).T)[0]
        ##print rs
        predict=rs.argmax()
        print "%d %d %f"%(predict,y_test[:,idx].argmax(),rs[predict])
        ax.set_title("%d"%(predict))
        plt.imshow(X_test[idx,0],cmap ='gray')
    plt.tight_layout()
    plt.show()